import java.io.IOException;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.CtNewMethod;
import javassist.NotFoundException;


public class JassistTiming {
	public static void main(String[] args) {
		if (args.length == 2){
			try {
				CtClass  cls = ClassPool.getDefault().get(args[0]);
				if (cls == null){
					System.err.println("Class " + args[0] + " not found");
				}else{
					addTiming(cls,args[1]);
					cls.writeFile();
					System.out.println("Added Timing to method " + args[0] + "."+ args[1]);
				}
				
			} catch (NotFoundException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} catch (IOException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} catch (CannotCompileException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		}else {
			System.out.println("Usage:JavassistTiming class method-name");
		}
	}

	private static void addTiming(CtClass cls, String mname) throws NotFoundException, CannotCompileException {
		CtMethod mold  = cls.getDeclaredMethod(mname);
		String nname = mname + "$impl";
		mold.setName(nname);
		CtMethod mnew = CtNewMethod.copy(mold, mname, cls, null);
		
		String type = mold.getReturnType().getName();
		StringBuffer body = new StringBuffer();
		body.append("{\nlong start = System.currentTimeMillis();\n");
		if (!"void".equals(type)){
			body.append(type+" result = ");
		}
		body.append(nname + "($$);\n");
		 body.append("System.out.println(\"Call to method " + mname +
		            " took \" +\n (System.currentTimeMillis()-start) + " +
		            "\" ms.\");\n");
		        if (!"void".equals(type)) {
		            body.append("return result;\n");
		        }
		        body.append("}");
		        
		mnew.setBody(body.toString());
		cls.addMethod(mnew);
		
		System.out.println("Interceptor method body:");
		System.out.println(body.toString());
	}
}
